#include "bzbshell/shell.h"
#include <sstream>

#if defined(WINDOWS)

#define WINDOWS_LEAN_AND_MEAN 1
#include <windows.h>

// __stdcall ignored here
using CheckTokenMembership_ptr = BOOL(*)(HANDLE TokenHandle, PSID SidToCheck, PBOOL IsMember);

#if defined(_COMPILER_MSVC)

#define POPEN_F _popen
#define PCLOSE_F _pclose

#else

#include <cstdio>
#define POPEN_F popen
#define PCLOSE_F pclose
// On some implementations, defining these functions is needed
// extern "C" FILE *popen(const char *command, const char *mode);
// extern "C" int pclose(FILE *stream);

#endif

#elif defined(LINUX)

#define POPEN_F popen
#define PCLOSE_F pclose

#include <unistd.h>
#include <sys/ioctl.h>
#include <sys/uio.h>
#include <linux/limits.h>
#include <signal.h>
#include <syscall.h>

#endif


std::string bzbshell::get_cwd() {

#if defined(WINDOWS)

	char buffer[256];
	int size = GetCurrentDirectory(256, buffer);
	if(!size) {
		bzbutil::throw_sys_error();
	}

	buffer[size] = '\0';

	return buffer;

#elif defined(LINUX)

	char buffer[PATH_MAX];
	char* res = getcwd(buffer, PATH_MAX);
	if(!res) {
		bzbutil::throw_sys_error();
	}

	return buffer;

#endif
}


void bzbshell::set_cwd(std::string path) {

#if defined(WINDOWS)

	int res = SetCurrentDirectory(path.c_str());
	if(!res) {
		bzbutil::throw_sys_error();
	}

#elif defined(LINUX)

	int res = chdir(path.c_str());
	if(res == -1) {
		bzbutil::throw_sys_error();
	}

#endif

}


bool bzbshell::is_privileged() {

#if defined(WINDOWS)

	SID_IDENTIFIER_AUTHORITY NtAuthority = SECURITY_NT_AUTHORITY;
	PSID AdministratorsGroup;
	BOOL b = AllocateAndInitializeSid(&NtAuthority, 2,
									 SECURITY_BUILTIN_DOMAIN_RID,
									 DOMAIN_ALIAS_RID_ADMINS,
									 0, 0, 0, 0, 0, 0,
									 &AdministratorsGroup);

	HINSTANCE lib = LoadLibraryEx("Advapi32.dll", nullptr, 0x00);

	if(!lib) {
		bzbutil::throw_sys_error();
	}

	CheckTokenMembership_ptr CheckTokenMembership =
		(CheckTokenMembership_ptr) GetProcAddress(lib, "CheckTokenMembership");

	FreeLibrary(lib);

	if(b) {
		if (!CheckTokenMembership(0, AdministratorsGroup, &b)) {
			b = FALSE;
		}

		FreeSid(AdministratorsGroup);
	}

	return b;

#elif defined(LINUX)

	return geteuid() == 0;

#endif
	
}


int bzbshell::pipe_open(const std::string &command) {

	FILE *open_pipe = POPEN_F(command.c_str(), "r");
	if (!open_pipe) {
		bzbutil::throw_sys_error();
	}

	char lines[BUFSIZ];
	while(fgets(lines, BUFSIZ, open_pipe)) {}

	errno = 0;
	int res = PCLOSE_F(open_pipe);

	if (res == -1 && errno) {
		bzbutil::throw_sys_error();
	}

	return res;
}


int bzbshell::pipe_open(const std::string &command, std::string &output) {

	FILE *open_pipe = POPEN_F(command.c_str(), "r");
	if (!open_pipe) {
		bzbutil::throw_sys_error(errno);
	}

	std::stringstream outstream;

	char lines[BUFSIZ];
	while(fgets(lines, BUFSIZ, open_pipe)) {
		outstream << lines;
	}

	output = outstream.str();

	errno = 0;
	int res = PCLOSE_F(open_pipe);

	if (res == -1 && errno) {
		bzbutil::throw_sys_error(errno);
	}

	return res;
}


#ifdef LINUX


void bzbshell::get_size(size_t *rows, size_t *columns) {

	struct winsize w;
	ioctl(STDOUT_FILENO, TIOCGWINSZ, &w);

	*rows = w.ws_row;
	*columns = w.ws_col;

}


#endif


#ifdef WINDOWS


void bzbshell::set_title(const std::string &title) {

	HWND window = GetConsoleWindow();
	int res = SetWindowText(window, title.c_str());
	if(!res) {
		bzbutil::throw_sys_error();
	}

}


void bzbshell::run_as_admin(std::string program) {

	HINSTANCE i = ShellExecute(nullptr,
								"runas",
								program.c_str(),
								nullptr,
								nullptr,
								SW_SHOWNORMAL);

	if((unsigned int) i != 0x2a) {
		bzbutil::throw_sys_error();
	}

}


#endif
